#include "vmlinux.h"
#include <bpf/bpf_helpers.h>


struct {
	__uint(type, BPF_MAP_TYPE_LPM_TRIE);
	__uint(max_entries, 16);
	__type(key, u64);
	__type(value, u32);
} blacklist SEC(".maps");

/// @ifindex 3
/// @flags 0
/// @xdpopts {"old_prog_fd":0}
SEC("xdp")
int xdp_pass(struct xdp_md* ctx) {
    void* data = (void*)(long)ctx->data;
    void* data_end = (void*)(long)ctx->data_end;
    int pkt_sz = data_end - data;

    // Only IPv4 supported for this example
    struct ethhdr *ether = data;
    if (data + sizeof(*ether) > data_end) {
        // Malformed Ethernet header
        return XDP_ABORTED;
    }

    if (ether->h_proto != 0x08U) {  // htons(ETH_P_IP) -> 0x08U
    // Non IPv4 traffic
    return XDP_PASS;
    }

    data += sizeof(*ether);
    struct iphdr *ip = data;
    if (data + sizeof(*ip) > data_end) {
    // Malformed IPv4 header
    return XDP_ABORTED;
    }

    struct {
        __u32 prefixlen;
        __u32 saddr;
    } key;

    key.prefixlen = 32;
    key.saddr = ip->saddr;

    // Lookup SRC IP in blacklisted IPs
    __u64 *rule_idx = bpf_map_lookup_elem(&blacklist, &key);
    if (rule_idx) {
        // Matched, increase match counter for matched "rule"
        __u32 index = *(__u32*)rule_idx;  // make verifier happy
        bpf_printk("[Drop] packet size is %d", pkt_sz);
        return XDP_DROP;
    }

    bpf_printk("[Pass] packet size is %d", pkt_sz);
    return XDP_PASS;
}

char __license[] SEC("license") = "GPL";